import sys
import atexit
import importlib
import os
import time
import warnings
import multiprocessing

warnings.filterwarnings('ignore', category=FutureWarning)

from task import test_scenario

def exit_handler():
    print("Exiting program...")

class NullWriter:
    def write(self, message):
        pass  # 不做任何操作

    def flush(self):
        pass  # 不做任何操作

def main():
    dir_current_file = os.path.dirname(__file__)
    dir_parent_1 = os.path.dirname(dir_current_file)
    dir_save_img = os.path.abspath(os.path.join(dir_parent_1, 'onsite_images_saved'))

    tic = time.time()

    # 加载场景
    scenarioOrganizer = importlib.import_module("onsite-unstructured.dynamic_scenes.scenarioOrganizer1")
    so = scenarioOrganizer.ScenarioOrganizer()
    dir_inputs = os.path.abspath(os.path.join(dir_parent_1, 'inputs'))
    dir_outputs = os.path.abspath(os.path.join(dir_parent_1, 'outputs'))
    so.load(dir_inputs, dir_outputs)

    # 准备测试场景列表
    scenario_list = []
    while True:
        scenario_to_test = so.next()
        if scenario_to_test is None:
            break
        scenario_list.append(scenario_to_test)

    # 设置最大同时运行的进程数
    max_concurrent_processes = 8
    processes = []
    scene_name_list = []

    for scenario_to_test in scenario_list:
        scene_name = scenario_to_test['data']['scene_name']
        # 检查当前正在运行的进程数量
        while len([p for p in processes if p.is_alive()]) >= max_concurrent_processes:
            time.sleep(1)

        p = multiprocessing.Process(target=test_scenario, args=(scenario_to_test, dir_save_img, dir_parent_1))
        p.start() if scene_name not in scene_name_list else print(f"场景 {scene_name} 已经在运行")
        scene_name_list.append(scene_name) if scene_name not in scene_name_list else None
        processes.append(p)

    # 等待所有进程完成
    for p in processes:
        p.join()

    toc = time.time()
    print(f"###log### 总用时: {int((toc - tic) // 3600)}小时 {int(((toc - tic) % 3600) // 60)}分钟 {(toc - tic) % 60:.2f}秒\n")



if __name__ == "__main__":
    # null_writer = NullWriter()
    # # 重定向标准输出
    # sys.stdout = null_writer
    multiprocessing.set_start_method('spawn', force=True)  # 根据系统设置启动方法
    atexit.register(exit_handler)
    main()
